#include "ntddk.h"
#include "AsmCall.h"
#include "Util.h"
#include "ntstrsafe.h"
#include "common.h"
#include "ia32_types.h"

PDEVICE_OBJECT pRealDeviceObj[26];
ULONG nCurrentHookDeviceObj;
PIRP pHookIrp;

NTSTATUS VmpAsmFindDeviceObjectStart(void* context) {
	ULONG32 szArrary[5];
	//Asm_cpuid('FDDO', szArrary);
	Asm_VmCall(HypercallNumber::kHookNtIofCallDriver, szArrary);
	return STATUS_SUCCESS;
}

NTSTATUS VmpAsmFindDeviceObjectEnd(void* context) {
	ULONG32 szArrary[5];
	//Asm_cpuid('UHFD', szArrary);
	Asm_VmCall(HypercallNumber::kUnHookIofCallDriver, szArrary);
	return STATUS_SUCCESS;
}

NTSTATUS SuperStripFileAttributesComplete(PDEVICE_OBJECT pDeviceObject, PIRP Irp, PVOID Context) {
	if (pHookIrp != 0)
	{
		if (Irp == pHookIrp) {
			KeSetEvent(Irp->UserEvent, IO_NO_INCREMENT, FALSE);
		}
	}
	else {
		KeSetEvent(Irp->UserEvent, IO_NO_INCREMENT, FALSE);
	}

	return STATUS_MORE_PROCESSING_REQUIRED;
}

BOOLEAN SuperStripFileAttributes(HANDLE hFile, PDEVICE_OBJECT pRealDevice) {
	PFILE_OBJECT pFileObject;

	NTSTATUS nStatus = ObReferenceObjectByHandle(hFile, DELETE, *IoFileObjectType, KernelMode, (PVOID *)&pFileObject, 0);
	if (!NT_SUCCESS(nStatus))
	{
		return FALSE;
	}

	PDEVICE_OBJECT pDeviceObject = pRealDevice ? pRealDevice : IoGetRelatedDeviceObject(pFileObject);
	if (pDeviceObject == 0)
	{
		ObDereferenceObject(pFileObject);
		return FALSE;
	}

	PIRP pIrp = IoAllocateIrp(pDeviceObject->StackSize, TRUE);
	if (pIrp == 0)
	{
		ObDereferenceObject(pFileObject);
		return FALSE;
	}

	KEVENT *pEvent = (KEVENT *)ExAllocatePoolWithTag(NonPagedPool, sizeof(KEVENT), kHyperPlatformCommonPoolTag);
	if (pEvent != 0)
	{
		KeInitializeEvent(pEvent, SynchronizationEvent, FALSE);

		FILE_BASIC_INFORMATION* pFileInformation = (FILE_BASIC_INFORMATION*)ExAllocatePoolWithTag(NonPagedPool, sizeof(FILE_BASIC_INFORMATION), kHyperPlatformCommonPoolTag);
		if (pFileInformation != 0)
		{
			memset(pFileInformation, 0, sizeof(FILE_BASIC_INFORMATION));

			(*pFileInformation).FileAttributes = FILE_ATTRIBUTE_NORMAL;

			pIrp->AssociatedIrp.SystemBuffer = pFileInformation;
			pIrp->UserEvent = pEvent;

			IO_STATUS_BLOCK ioStatusBlock;
			pIrp->UserIosb = &ioStatusBlock;
			pIrp->Tail.Overlay.OriginalFileObject = pFileObject;
			pIrp->Tail.Overlay.Thread = (PETHREAD)KeGetCurrentThread();
			pIrp->RequestorMode = KernelMode;

			PIO_STACK_LOCATION pIoStackLocation = IoGetNextIrpStackLocation(pIrp);
			pIoStackLocation->MajorFunction = IRP_MJ_SET_INFORMATION;
			pIoStackLocation->DeviceObject = pDeviceObject;
			pIoStackLocation->FileObject = pFileObject;
			pIoStackLocation->Parameters.SetFile.Length = sizeof(FILE_BASIC_INFORMATION);
			pIoStackLocation->Parameters.SetFile.FileInformationClass = FileBasicInformation;
			pIoStackLocation->Parameters.SetFile.FileObject = pFileObject;

			pHookIrp = pIrp;

			IoSetCompletionRoutine(pIrp, SuperStripFileAttributesComplete, 0, TRUE, TRUE, TRUE);
			IoCallDriver(pDeviceObject, pIrp);
			KeWaitForSingleObject(pEvent, Executive, KernelMode, TRUE, NULL);

			pHookIrp = 0;

			ExFreePoolWithTag(pFileInformation, kHyperPlatformCommonPoolTag);
		}
		ExFreePoolWithTag(pEvent, kHyperPlatformCommonPoolTag);
	}
	
	ObDereferenceObject(pFileObject);

	IoFreeIrp(pIrp);

	return TRUE;
}

BOOLEAN SuperDeleteFile(HANDLE hFile, PDEVICE_OBJECT pRealDevice) {
	
	SuperStripFileAttributes(hFile,pRealDevice);
	PFILE_OBJECT pFileObject;
	NTSTATUS nStatus = ObReferenceObjectByHandle(hFile, DELETE, *IoFileObjectType, KernelMode, (PVOID*)&pFileObject, 0);
	if (!NT_SUCCESS(nStatus))
	{
		return FALSE;
	}

	PDEVICE_OBJECT pDeviceObject = pRealDevice ? pRealDevice : IoGetRelatedDeviceObject(pFileObject);
	if (pDeviceObject == 0)
	{
		ObDereferenceObject(pFileObject);
		return FALSE;
	}

	PIRP pIrp = IoAllocateIrp(pDeviceObject->StackSize, TRUE);
	if (pIrp == 0)
	{
		ObDereferenceObject(pFileObject);
		return FALSE;
	}

	KEVENT event;
	KeInitializeEvent(&event, SynchronizationEvent, FALSE);

	FILE_DISPOSITION_INFORMATION fileDispositionInfor;
	fileDispositionInfor.DeleteFile = TRUE;
	
	pIrp->AssociatedIrp.SystemBuffer = &fileDispositionInfor;
	pIrp->UserEvent = &event;
	
	IO_STATUS_BLOCK ioStatusBlock;
	pIrp->UserIosb = &ioStatusBlock;
	pIrp->Tail.Overlay.OriginalFileObject = pFileObject;
	pIrp->Tail.Overlay.Thread = (PETHREAD)KeGetCurrentThread();
	pIrp->RequestorMode = KernelMode;

	PIO_STACK_LOCATION pIoStackLocation = IoGetNextIrpStackLocation(pIrp);
	pIoStackLocation->MajorFunction = IRP_MJ_SET_INFORMATION;
	pIoStackLocation->DeviceObject = pDeviceObject;
	pIoStackLocation->FileObject = pFileObject;
	pIoStackLocation->Parameters.SetFile.Length = sizeof(FILE_DISPOSITION_INFORMATION);
	pIoStackLocation->Parameters.SetFile.FileInformationClass = FileDispositionInformation;
	pIoStackLocation->Parameters.SetFile.FileObject = pFileObject;

	IoSetCompletionRoutine(pIrp, SuperStripFileAttributesComplete, 0, TRUE, TRUE, TRUE);

	PSECTION_OBJECT_POINTERS pSectionObjectPointer = pFileObject->SectionObjectPointer;
	pSectionObjectPointer->ImageSectionObject = 0;
	pSectionObjectPointer->DataSectionObject = 0;

	IoCallDriver(pDeviceObject, pIrp);

	KeWaitForSingleObject(&event, Executive, KernelMode, TRUE, NULL);
	ObDereferenceObject(pFileObject);

	IoFreeIrp(pIrp);

	return TRUE;
}

HANDLE SuperIoOpenFile(
	IN PCWSTR FileName,
	IN ACCESS_MASK DesiredAccess,
	IN ULONG ShareAccess,
	ULONG Disposition
)
{
	NTSTATUS              ntStatus;
	UNICODE_STRING        uniFileName;
	OBJECT_ATTRIBUTES     objectAttributes;
	HANDLE                ntFileHandle;
	IO_STATUS_BLOCK       ioStatus;

	RtlInitUnicodeString(&uniFileName, FileName);
	InitializeObjectAttributes(&objectAttributes, &uniFileName,
		OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, NULL);

	ntStatus = IoCreateFile(&ntFileHandle,
		DesiredAccess,
		&objectAttributes,
		&ioStatus,
		0,
		FILE_ATTRIBUTE_NORMAL,
		ShareAccess,
		Disposition,
		0,
		NULL,
		0,
		(CREATE_FILE_TYPE)0,
		NULL,
		IO_NO_PARAMETER_CHECKING);

	if (!NT_SUCCESS(ntStatus))
	{
		return 0;
	}

	return ntFileHandle;
}

BOOLEAN SuperTryToFindDeviceByDeleteFile(HANDLE hFile, PDEVICE_OBJECT pRealDevice) {
	UtilForEachProcessor(VmpAsmFindDeviceObjectStart, 0);
	SuperStripFileAttributes(hFile, pRealDevice);
	UtilForEachProcessor(VmpAsmFindDeviceObjectEnd, 0);
	PFILE_OBJECT pFileObject;
	NTSTATUS nStatus = ObReferenceObjectByHandle(hFile, DELETE, *IoFileObjectType, KernelMode, (PVOID*)&pFileObject, 0);
	if (!NT_SUCCESS(nStatus))
	{
		return FALSE;
	}

	PDEVICE_OBJECT pDeviceObject = pRealDevice ? pRealDevice : IoGetRelatedDeviceObject(pFileObject);
	if (pDeviceObject == 0)
	{
		ObDereferenceObject(pFileObject);
		return FALSE;
	}

	PIRP pIrp = IoAllocateIrp(pDeviceObject->StackSize, TRUE);
	if (pIrp == 0)
	{
		ObDereferenceObject(pFileObject);
		return FALSE;
	}

	KEVENT event;
	KeInitializeEvent(&event, SynchronizationEvent, FALSE);

	FILE_DISPOSITION_INFORMATION fileDispositionInfor;
	fileDispositionInfor.DeleteFile = TRUE;

	pIrp->AssociatedIrp.SystemBuffer = &fileDispositionInfor;
	pIrp->UserEvent = &event;

	IO_STATUS_BLOCK ioStatusBlock;
	pIrp->UserIosb = &ioStatusBlock;
	pIrp->Tail.Overlay.OriginalFileObject = pFileObject;
	pIrp->Tail.Overlay.Thread = (PETHREAD)KeGetCurrentThread();
	pIrp->RequestorMode = KernelMode;

	PIO_STACK_LOCATION pIoStackLocation = IoGetNextIrpStackLocation(pIrp);
	pIoStackLocation->MajorFunction = IRP_MJ_SET_INFORMATION;
	pIoStackLocation->DeviceObject = pDeviceObject;
	pIoStackLocation->FileObject = pFileObject;
	pIoStackLocation->Parameters.SetFile.Length = sizeof(FILE_DISPOSITION_INFORMATION);
	pIoStackLocation->Parameters.SetFile.FileInformationClass = FileDispositionInformation;
	pIoStackLocation->Parameters.SetFile.FileObject = pFileObject;

	IoSetCompletionRoutine(pIrp, SuperStripFileAttributesComplete, 0, TRUE, TRUE, TRUE);

	PSECTION_OBJECT_POINTERS pSectionObjectPointer = pFileObject->SectionObjectPointer;
	pSectionObjectPointer->ImageSectionObject = 0;
	pSectionObjectPointer->DataSectionObject = 0;

	IoCallDriver(pDeviceObject, pIrp);

	KeWaitForSingleObject(&event, Executive, KernelMode, TRUE, NULL);
	ObDereferenceObject(pFileObject);

	IoFreeIrp(pIrp);

	return TRUE;
}

void FindAllDriverRealDevice() {
	pHookIrp = 0;
	nCurrentHookDeviceObj = 0;
	for (size_t i = 0; i < 26; i++)
	{
		pRealDeviceObj[i] = 0;
		WCHAR pDriverPath[1024];
		UNICODE_STRING dtUnicodeString;
		RtlStringCbPrintfW((NTSTRSAFE_PWSTR)pDriverPath, sizeof(pDriverPath), L"\\??\\%C:\\whisper_test_create.file.test", 'A' + i);
		RtlInitUnicodeString(&dtUnicodeString, (PCWSTR)pDriverPath);
		HANDLE hFileHandle = SuperIoOpenFile((PCWSTR)pDriverPath, FILE_READ_ACCESS, FILE_SHARE_READ, FILE_OPEN_IF);
		if (hFileHandle != 0)
		{
			SuperTryToFindDeviceByDeleteFile(hFileHandle, 0);			
			ZwClose(hFileHandle);
		}

		pHookIrp = 0;
		nCurrentHookDeviceObj++;
	}

	for (size_t i = 0; i < 26; i++)
	{
		KdPrint(("%C  ---------> 0x%p\n", 'A' + i, pRealDeviceObj[i]));
	}
}

BOOLEAN SuperKillFile(PCWSTR pFilePath) {
	if (pFilePath != NULL)
	{
		HANDLE hFileHandle = SuperIoOpenFile((PCWSTR)pFilePath, FILE_READ_ACCESS, FILE_SHARE_READ, FILE_OPEN);
		if (hFileHandle != 0)
		{
			BOOLEAN bReturn = FALSE;
			if (pFilePath[0] == '\\' && pFilePath[1] == '?' && pFilePath[2] == '?' && pFilePath[3] == '\\')
			{
				if (pFilePath[4] >= 'A' && pFilePath[4] <= 'Z')
				{
					if (pRealDeviceObj[pFilePath[4] - 'A'] == 0)
					{
						nCurrentHookDeviceObj = pFilePath[4] - 'A';
						WCHAR pDriverPath[1024];
						UNICODE_STRING dtUnicodeString;
						RtlStringCbPrintfW((NTSTRSAFE_PWSTR)pDriverPath, sizeof(pDriverPath), L"\\??\\%C:\\whisper_test_create.file.test", pFilePath[4]);
						RtlInitUnicodeString(&dtUnicodeString, (PCWSTR)pDriverPath);
						HANDLE hTestFileHandle = SuperIoOpenFile((PCWSTR)pDriverPath, FILE_READ_ACCESS, FILE_SHARE_READ, FILE_OPEN_IF);
						if (hTestFileHandle != 0)
						{
							KdPrint(("[Whisper]TestFindHandle: %p\n", hTestFileHandle));
							SuperTryToFindDeviceByDeleteFile(hTestFileHandle, 0);
							ZwClose(hTestFileHandle);
						}

						pHookIrp = 0;
					}

					if (pRealDeviceObj[pFilePath[4] - 'A'] != 0) {
						SuperDeleteFile(hFileHandle, pRealDeviceObj[pFilePath[4] - 'A']);
						bReturn = TRUE;
					}
				}
				else if (pFilePath[4] >= 'a' && pFilePath[4] <= 'z') {
					if (pRealDeviceObj[pFilePath[4] - 'a'] == 0)
					{
						nCurrentHookDeviceObj = pFilePath[4] - 'a';
						WCHAR pDriverPath[1024];
						UNICODE_STRING dtUnicodeString;
						RtlStringCbPrintfW((NTSTRSAFE_PWSTR)pDriverPath, sizeof(pDriverPath), L"\\??\\%C:\\whisper_test_create.file.test", pFilePath[4]);
						RtlInitUnicodeString(&dtUnicodeString, (PCWSTR)pDriverPath);
						HANDLE hTestFileHandle = SuperIoOpenFile((PCWSTR)pDriverPath, FILE_READ_ACCESS, FILE_SHARE_READ, FILE_OPEN_IF);
						if (hTestFileHandle != 0)
						{
							SuperTryToFindDeviceByDeleteFile(hTestFileHandle, 0);
							ZwClose(hTestFileHandle);
						}

						pHookIrp = 0;
					}
					
					if (pRealDeviceObj[pFilePath[4] - 'a'] != 0) {
						SuperDeleteFile(hFileHandle, pRealDeviceObj[pFilePath[4] - 'a']);
						bReturn = TRUE;
					}
				}
			}
			ZwClose(hFileHandle);
			return bReturn;
		}
		else {
			return FALSE;
		}
	}
	else {
		return FALSE;
	}
}